from pydantic import BaseModel, Field, validator, field_validator
from typing import Optional, Any, Dict, List
from enum import Enum
from datetime import datetime, timedelta

# Enums for validation
class DataType(str, Enum):
    IMAGE = "01"
    TIMESERIES = "02"
    TEXT = "10"

class StrategyType(str, Enum):
    EDGE_FIRST = "00"
    CLOUD_FIRST = "01"
    COLLABORATIVE = "10"

class TaskType(str, Enum):
    CLOUD_EVOLUTION = "01"
    COLLABORATIVE_EVOLUTION = "02"

class TaskStatus(str, Enum):
    NOT_STARTED = "NOT_STARTED"
    RUNNING = "RUNNING"
    COMPLETED = "COMPLETED"
    TERMINATED = "TERMINATED"
    ERROR = "ERROR"

class ModelType(str, Enum):
    VIT = "Vit"
    LLAMA = "Llama"
    LLAVA = "Llava"
    INFORMER = "Informer"
    CLIP = "CLIP"
    RESNET = "Resnet"
    VSE = "VSE"

class OptimizationMethod(str, Enum):
    SGD = "SGD"
    QUEST = "Quest"
    LORA = "Lora"
    INFONCE = "InfoNCE"
    DISTILL = "distill"
    FEDAVG = "FedAVG"
    QUANTI = "Quanti"

# Request Models
class UpdateSamplesRequest(BaseModel):
    edgeID: List[int]
    threshold: float = Field(..., ge=0, le=1)
    SCloud: bool = True
    ZCloud: bool = True

class CouldEvolutionRequest(BaseModel):
    model: ModelType
    task: str
    method: OptimizationMethod = OptimizationMethod.QUEST
    lr: float = 0.001
    epoch: int = Field(5, ge=1, le=10)
    batchSize: int = 32
    dataStartTime: str
    dataEndTime: str

    # 枚举类型转换
    @field_validator('model', mode='before')
    @classmethod
    def validate_model(cls, v):
        if isinstance(v, str):
            try:
                return ModelType(v.upper())
            except ValueError:
                return ModelType.VIT
        return v

    @field_validator('method', mode='before')
    @classmethod
    def validate_method(cls, v):
        if isinstance(v, str):
            try:
                return OptimizationMethod(v.upper())
            except ValueError:
                return OptimizationMethod.QUEST
        return v

    # 浮点数转换
    @field_validator('lr', mode='before')
    @classmethod
    def validate_lr(cls, v):
        if isinstance(v, str):
            try:
                return float(v)
            except ValueError:
                return 0.001
        return v

    # 整数转换
    @field_validator('epoch', mode='before')
    @classmethod
    def validate_epoch(cls, v):
        if isinstance(v, str):
            try:
                value = int(float(v))
                if value < 1:
                    return 1
                elif value > 10:
                    return 10
                return value
            except ValueError:
                return 5
        return v

    @field_validator('batchSize', mode='before')
    @classmethod
    def validate_batch_size(cls, v):
        if isinstance(v, str):
            try:
                return int(float(v))
            except ValueError:
                return 32
        return v

    # 日期时间转换
    @field_validator('dataStartTime', 'dataEndTime', mode='before')
    @classmethod
    def validate_dates(cls, v, info):
        if isinstance(v, str):
            try:
                datetime.fromisoformat(v)
                return v
            except ValueError:
                try:
                    dt = datetime.strptime(v, '%Y-%m-%d')
                    return dt.isoformat()
                except ValueError:
                    # 根据字段名返回默认值
                    if info.field_name == 'dataStartTime':
                        return datetime.now().isoformat()
                    else:  # dataEndTime
                        return (datetime.now() + timedelta(days=1)).isoformat()
        return v

    class Config:
        extra = "ignore"
        validate_assignment = True

class CollaborateEvolutionRequest(BaseModel):
    Smodel: ModelType
    Zmodel: ModelType
    task: str
    method: OptimizationMethod = OptimizationMethod.QUEST
    lr: float = 0.001
    epoch: int = Field(5, ge=1, le=10)
    batchSize: int = 32
    dataStartTime: str
    dataEndTime: str

    # 枚举类型转换
    @field_validator('Smodel', 'Zmodel', mode='before')
    @classmethod
    def validate_models(cls, v):
        if isinstance(v, str):
            try:
                return ModelType(v.upper())
            except ValueError:
                return ModelType.VIT
        return v

    @field_validator('method', mode='before')
    @classmethod
    def validate_method(cls, v):
        if isinstance(v, str):
            try:
                return OptimizationMethod(v.upper())
            except ValueError:
                return OptimizationMethod.QUEST
        return v

    # 浮点数转换
    @field_validator('lr', mode='before')
    @classmethod
    def validate_lr(cls, v):
        if isinstance(v, str):
            try:
                return float(v)
            except ValueError:
                return 0.001
        return v

    # 整数转换
    @field_validator('epoch', mode='before')
    @classmethod
    def validate_epoch(cls, v):
        if isinstance(v, str):
            try:
                value = int(float(v))
                if value < 1:
                    return 1
                elif value > 10:
                    return 10
                return value
            except ValueError:
                return 5
        return v

    @field_validator('batchSize', mode='before')
    @classmethod
    def validate_batch_size(cls, v):
        if isinstance(v, str):
            try:
                return int(float(v))
            except ValueError:
                return 32
        return v

    # 日期时间转换
    @field_validator('dataStartTime', 'dataEndTime', mode='before')
    @classmethod
    def validate_dates(cls, v, info):
        if isinstance(v, str):
            try:
                datetime.fromisoformat(v)
                return v
            except ValueError:
                try:
                    dt = datetime.strptime(v, '%Y-%m-%d')
                    return dt.isoformat()
                except ValueError:
                    if info.field_name == 'dataStartTime':
                        return datetime.now().isoformat()
                    else:  # dataEndTime
                        return (datetime.now() + timedelta(days=1)).isoformat()
        return v

    class Config:
        extra = "ignore"
        validate_assignment = True

class TaskProgressRequest(BaseModel):
    taskID: str

class StrategyRequest(BaseModel):
    dataType: DataType
    strategyType: StrategyType

class InferenceRequest(BaseModel):
    filePath: str
    dataType: DataType = DataType.IMAGE

class BatchInferenceRequest(BaseModel):
    file: str
    dataType: DataType = DataType.IMAGE

class GPUStatusRequest(BaseModel):
    pass

class CPUStatusRequest(BaseModel):
    pass

class NetworkStatusRequest(BaseModel):
    pass

class GetFileRequest(BaseModel):
    filePath: str

# Response Models
class UpdateSamplesResponse(BaseModel):
    num: int
    sampleID: List[str]

class TaskResponse(BaseModel):
    taskID: str
    taskType: TaskType
    status: TaskStatus
    sampleID: List[str]

class TaskProgressResponse(BaseModel):
    taskID: str
    strategyType: TaskType
    status: TaskStatus
    loss: List[float]
    acc: List[float]
    rate: float
    time: str

class StrategyResponse(BaseModel):
    strategyType: StrategyType

class InferenceResponse(BaseModel):
    resultFilePath: str
    time: str

class BatchInferenceResponse(BaseModel):
    resultFilePath: List[str]
    time: str

class GPUStatusResponse(BaseModel):
    data: Dict[str, List[float]]

class CPUStatusResponse(BaseModel):
    data: Dict[str, List[float]]

class NetworkStatusResponse(BaseModel):
    data: Dict[str, List[float]]

class FileResponse(BaseModel):
    fileName: str
    fileSize: int
    fileContent: str

class BaseResponse(BaseModel):
    status: str = "success"
    message: str
    data: Optional[Any] = None
